load("//bazel:api.bzl", "modular_py_binary", "modular_py_test", "pkg_files", "requirement")

modular_py_binary(
    name = "kbench",
    srcs = [
        "__init__.py",
        "kbench.py",
        "utils.py",
    ],
    data = [
        "@nvshmem_prebuilt//:device",
        "@nvshmem_prebuilt//:host",
    ],
    env = {
        "KERNEL_BENCHMARKS_ROOT": "$$BUILD_WORKSPACE_DIRECTORY/" + package_name() + "/..",
        "BITCODE_LIBS_PATH": "$(rootpath @nvshmem_prebuilt//:device)",
        "MODULAR_SHMEM_LIB_DIR": "../+http_archive+nvshmem_prebuilt",
    },
    imports = ["."],
    main = "kbench.py",
    mojo_deps = [
        "//max:_cublas",
        "//max:_cudnn",
        "//max:_cufft",
        "//max:_curand",
        "//max:_rocblas",
        "//max:comm",
        "//max:internal_utils",
        "//max:kv_cache",
        "//max:layout",
        "//max:linalg",
        "//max:nn",
        "//max:quantization",
        "//max:register",
        "//max:shmem",
        "@mojo//:stdlib",
    ],
    deps = [
        requirement("click"),
        requirement("numpy"),
        requirement("pandas"),
        requirement("pygments"),
        requirement("pyyaml"),
        requirement("rich"),
        requirement("scipy"),
        requirement("plotly"),
        requirement("tabulate"),
    ],
)

modular_py_binary(
    name = "kprofile",
    srcs = [
        "kprofile.py",
    ],
    imports = ["."],
    main = "kprofile.py",
    mojo_deps = [
        "//max:internal_utils",
        "//max:kv_cache",
        "//max:layout",
        "//max:linalg",
        "//max:nn",
        "//max:quantization",
        "//max:register",
        "@mojo//:stdlib",
    ],
    deps = [
        ":kbench",
        requirement("rich"),
    ],
)

modular_py_binary(
    name = "kplot",
    srcs = [
        "kplot.py",
    ],
    imports = ["."],
    main = "kplot.py",
    mojo_deps = [
        "//max:internal_utils",
        "//max:kv_cache",
        "//max:layout",
        "//max:linalg",
        "//max:nn",
        "//max:quantization",
        "//max:register",
        "@mojo//:stdlib",
    ],
    deps = [
        ":kbench",
        ":kprofile",
        requirement("plotly"),
        requirement("kaleido"),
    ],
)

modular_py_binary(
    name = "kdiff",
    srcs = [
        "kdiff.py",
    ],
    imports = ["."],
    main = "kdiff.py",
    deps = [
        requirement("click"),
        requirement("pyyaml"),
    ],
)

modular_py_test(
    name = "autotune_tests",
    srcs = glob(["tests/*.py"]),
    data = glob(["tests/*.csv"]) + [
        "sample.mojo",
        "test.yaml",
    ],
    env = {"KERNEL_BENCHMARKS_ROOT": "max/kernels/benchmarks/"},
    mojo_deps = [
        "//max:internal_utils",
        "//max:kv_cache",
        "//max:layout",
        "//max:linalg",
        "//max:nn",
        "//max:quantization",
        "//max:register",
        "@mojo//:stdlib",
    ],
    deps = [
        ":kbench",
        ":kplot",
        ":kprofile",
    ],
)

pkg_files(
    name = "autotune_files",
    srcs = [
        "kbench.py",
        "kplot.py",
        "kprofile.py",
        "pyproject.toml",
        "utils.py",
    ],
    prefix = "kernel-benchmark/autotune",
    visibility = ["//visibility:public"],
)
